# This Python script is used to generate the flows
# to test for python simulation framework
import math
import os

class genTraverseFlows():
    def __init__(self, gpuNum_: int, maxleafNum_: int) -> None:
        self.gpuNum: int = gpuNum_
        self.maxleafNum: int = maxleafNum_
        self.testCases: list[list[int]] = []
    
    def genTestCase(self):
        for numOutLeaf in range(1, self.maxleafNum):
            for numInFlow in range(numOutLeaf, self.maxleafNum):
                # Too many in flows, continue
                self.testCases.append([])
                if numInFlow > numOutLeaf * self.gpuNum:
                    continue
                # How many out flows to a leaf switch
                numFlowtoLeaf = math.floor(numInFlow/numOutLeaf)
                assert(numFlowtoLeaf <= self.gpuNum)
                remainInFlow = numInFlow
                for i in range(numOutLeaf - 1):
                    # Append the num of flows send to the leaf switch to a list
                    self.testCases[-1].append(numFlowtoLeaf)
                    remainInFlow -= numFlowtoLeaf
                self.testCases[-1].append(remainInFlow)
        print('Generate {} testcases'.format(len(self.testCases)))
        return len(self.testCases)
        pass
    
    # Generate flows to test in omnet
    # write out to txt at the same time
    def genFlow(self):
        testCaseNum = 0
        for testCase in self.testCases:
            numSendLeaf = 0 # Send from which leaf switch
            numSentFlow = numSendLeaf
            filePath = './testFlows/testcase{}.txt'.format(str(testCaseNum + 1))
            with open(filePath, 'w', encoding='utf-8') as file:
                file.write('// This is testcase {}\n'.format(str(testCaseNum + 1)))
                file.write('// There are {} flows on a link\n'.format(testCase[-1]))
                file.write('levelname0=gpu\n')
                file.write('levelname1=nic\n')
                file.write('levelname2=Leaf_Switch\n')
                file.write('levelname3=Spine_Switch\n')
                file.write('datasize=10000000\n')
                file.write('time=0\n')
                leafIndex = 0
                for downLeafFLowNum in testCase:
                    for downFlowonLeaf in range(downLeafFLowNum):
                        file.write('// Flow {}\n'.format(numSentFlow + 1))
                        usedGpu = (self.maxleafNum - numSendLeaf - 1) * self.gpuNum + self.gpuNum - 1
                        file.write('compound=(0,{},0)\n'.format(str(usedGpu)))
                        file.write('compound=(1,{},0)\n'.format(str(usedGpu)))
                        file.write('compound=(2,{},0)\n'.format(str(self.maxleafNum - numSendLeaf - 1)))
                        file.write('compound=(3,0,0)\n') # Every flow flows through spine switch 0
                        file.write('compound=(2,{},0)\n'.format(str(leafIndex)))
                        file.write('compound=(1,{},0)\n'.format(str(leafIndex * self.gpuNum + downFlowonLeaf)))
                        file.write('compound=(0,{},0)\n'.format(str(leafIndex * self.gpuNum + downFlowonLeaf)))
                        file.write('\n')
                        numSentFlow += 1
                    leafIndex += 1
            testCaseNum += 1

def genIniFile(filePath, testCaseNum):
    with open(filePath, 'w', encoding='utf-8') as file:
        file.write('[General]\n')
        file.write('include ./general.ini\n')
        file.write('*.numMachine = 8\n')
        file.write('**.nic[*].typename = "EthernetNicSimple"\n')
        for i in range(testCaseNum):
            file.write('[Config testcase{}]\n'.format(str(i)))
            file.write('**.inputType = 1\n')
            file.write('**.file_path = "./testFlows/testcase{}.txt"\n\n'.format(i + 1))

if __name__ == '__main__':
    genFlow = genTraverseFlows(8, 8)
    testcasesnum = genFlow.genTestCase()
    genIniFile("simulations/roce_sim/test_flow.ini", testcasesnum)
    genFlow.genFlow()
    # Run a omnet simulation
    for i in range(testcasesnum):
        file = open("resultFile.txt", 'a', encoding='utf-8')
        file.write("\ntestcase{}\n".format(i))
        file.close()
        os.system('./ai-cluster-simulation -f ./simulations/roce_sim/test_flow.ini -u Cmdenv --cmdenv-express-mode=true -c testcase{}'.format(str(i)))
        # break
